import numpy as np
import cPickle
import os
import os.path as osp
from PIL import Image
import xml.etree.ElementTree as ET

def read_dataLst(dataLstPath):
    """
    read txt file to get imagename
    for some class, find imagename which consisting of this object
    """
    img_name_list   = []
    with open(dataLstPath, 'r') as fid:
        lines   = fid.readlines()
        for ele in lines:
            line= ele.strip()
            line_split  = line.split(" ")
            if len(line_split) == 1:
                img_name_list.append(line.strip())
            else:
                img_name, exist_obj = line_split[0], line_split[1]
                if exist_obj == '1':
                    img_name_list.append(img_name.strip())
                elif exist_obj == "-1":
                    pass
                else:
                    assert True, "the 2th col:{} is not 1 or -1".format(exist_obj)
    return img_name_list

def parse_voc_rec(filename):
    # refer to https://github.com/facebookresearch/Detectron/blob/master/detectron/datasets/voc_eval.py
    """Parse a PASCAL VOC xml file."""
    tree = ET.parse(filename)
    objects = []
    for obj in tree.findall('object'):
        obj_struct = {}
        obj_struct['name'] = obj.find('name').text
        obj_struct['pose'] = obj.find('pose').text
        obj_struct['truncated'] = int(obj.find('truncated').text)
        obj_struct['difficult'] = int(obj.find('difficult').text)
        bbox = obj.find('bndbox')
        obj_struct['bbox'] = [int(bbox.find('xmin').text),
                              int(bbox.find('ymin').text),
                              int(bbox.find('xmax').text),
                              int(bbox.find('ymax').text)]
        objects.append(obj_struct)
    return objects

def process_voc_rec(filename, clsToInd):
    """
    read xml data, and return obj and cls data
    """
    obj_list    = parse_voc_rec(filename)
    obj_num     = len(obj_list)

    bbox_arr    = np.zeros((obj_num, 4), dtype=np.float32)
    cls_arr     = np.zeros((obj_num,), dtype=np.int32)
    
    for i, obj_struct in enumerate(obj_list):
        isDiff  = obj_struct['difficult']
        bbox    = obj_struct['bbox']
        cls_ind = clsToInd[obj_struct['name']]
        if isDiff:
            # difficult sample need to be filter out
            cls_ind     = -1
        bbox_arr[i]     = np.array(bbox)
        cls_arr[i]      = int(cls_ind)
    return bbox_arr, cls_arr

def load_voc_annotation(img_dir, anno_dir, dataLstPath, data_dir='', img_ext='jpg'):
    """
    load dota annotation from xml file
    """
    assert osp.exists(img_dir), "img_dir: {} must exist".format(img_dir)
    assert osp.exists(anno_dir),"anno_dir: {} must exist".format(anno_dir)
    assert osp.exists(dataLstPath), "dataLstPath: {} must exist".format(dataLstPath)

    cats    = ['person', #Person
                'bird', 'cat', 'cow', 'dog', 'horse', 'sheep', #Animal
                'aeroplane', 'bicycle', 'boat', 'bus', 'car', 'motorbike', 'train', #Vehicle
                'bottle', 'chair', 'diningtable', 'pottedplant', 'sofa', 'tvmonitor']# Indoor
    classes = ['__background__'] + cats
    class_to_ind    = dict(zip(classes, range(len(classes))))

    roidb           = []

    # read DataLstPath
    data_name_list   = read_dataLst(dataLstPath)
    data_num         = len(data_name_list)

    noObj_list      = []
    for ind, data_name in enumerate(data_name_list):
        img_idx     = data_name
        img_name    = data_name + '.' + img_ext
        #print("{}-th: {}".format(ind, img_name))
        if ind % 500 == 0:
            print('{}/{} annotation file processed'.format(ind, data_num))
        img_path    = osp.join(img_dir, img_name)
        anno_name   = img_idx + '.xml'
        anno_path   = osp.join(anno_dir, anno_name)
        assert osp.exists(anno_path), "anno_name: {} not exist".format(anno_name)

        # get image info
        im  = Image.open(img_path)
        im_width, im_height = im.size
        bbox_arr, gt_classes= process_voc_rec(anno_path, class_to_ind)
        

        roi_rec = {'image': data_dir + img_name,
                   'height': im_height,
                   'width': im_width,
                   'boxes': bbox_arr,#(xmin, ymin, xmax, ymax)
                   'gt_classes': gt_classes,
                    }
        roidb.append(roi_rec)
    return roidb

if __name__ == '__main__':

    # old address
    # dataset_dir     = '/home/users/zewen.he/projects/dataset/DOTA-dataset'
    # datatype        = 'val'
    # img_dir         = osp.join(dataset_dir, datatype, 'images', 'images')
    # anno_dir        = osp.join(dataset_dir, datatype, 'labelTxt')
    # anno_relabel_dir= osp.join(anno_dir, datatype+'set_reclabelTxt')

    import socket
    hostname        = socket.gethostname()
    if 'mac' in hostname:
        dataset_dir     = '/Users/hezewen/projects/dataset/VOCdevkit/'
    elif 'gz-gpu' in hostname:
        dataset_dir     = '/home/users/zewen.he/projects/dataset/VOCdevkit/'
    elif 'tzz-PC' in hostname:
        dataset_dir     = '/home/tzz/projects/dataset/VOCdevkit/'
    elif 'zws-625' in hostname:
        dataset_dir     = '/home/hzw/projects/dataset/VOCdevkit/'
    else:
        assert False, 'hostname:{} not consists of VOCdevkit'.format(hostname)

    datatypeList    = ['train', 'val', 'trainval']
    imagesetList    = ['2007', '2012']

    for i in range(len(datatypeList)):
        for j in range(len(imagesetList)):
            datatype= datatypeList[i]
            imageset= imagesetList[j]
            if datatype == 'test' and imageset == '2012':
                continue
#    datatype        = 'train'
#    imageset        = "VOC2012"
            print("Proessing datatype: {}, imageset:{}".format(datatype, imageset))
            img_dir         = osp.join(dataset_dir, 'VOC'+imageset, 'JPEGImages')
            anno_dir        = osp.join(dataset_dir, 'VOC'+imageset, 'Annotations')
            dataLst_path    = osp.join(dataset_dir, 'VOC'+imageset, 'ImageSets/Main/{}.txt'.format(datatype))

            gt_roidb        = load_voc_annotation(img_dir, anno_dir, dataLst_path, data_dir=img_dir+'/')

            task            = 'det'
            roidb_dir       = osp.join(dataset_dir, 'roidbs')
            if osp.exists(roidb_dir) is False:
                os.makedirs(roidb_dir)
                print("mkdir roidb_dir: {}".format(roidb_dir))
            save_roidb_path = osp.join(roidb_dir, '{}{}_{}_gt_roidb.pkl'.format(datatype, imageset, task))
            print("begin write to disk, save_roidb_path: {}".format(save_roidb_path))
            with open(save_roidb_path, 'wb') as fid:
                cPickle.dump(gt_roidb, fid, cPickle.HIGHEST_PROTOCOL)
            print("end write to disk")
